from .relys import BitSet
from .relys import Chip

class Memory(Chip):
    def __init__(self, awidth, dwidth):
        super().__init__(awidth+dwidth*2+3, [0,awidth+dwidth+3], [awidth+dwidth+3,awidth+dwidth*2+3])
        self.AI = self.pinIN@[0]
        self.RI = self.pinIN@[1]
        self.RO = self.pinIN@[2]
        self.ADDR = self.pinIN@[3, awidth+3]
        self.DAT = self.pinIN@[awidth+3, awidth+dwidth+3]
        self.width = dwidth
        self.mem = BitSet(dwidth*pow(2,awidth))
        self._dat = self.read_mem(0)
    def read_mem(self, addr):
        return self.mem@[addr*self.width, (addr+1)*self.width]
    def prog_mem(self, addr, data):
        self.mem[addr*self.width:(addr+1)*self.width] = data

@Memory.keyframe(1)
def _frame(chip):
    if chip.RO() : chip.pinOUT(chip._dat())
    else : chip.pinOUT(0)
    return chip

@Memory.keyframe(2)
def _frame(chip):
    if chip.AI() == 1 : chip._dat = chip.read_mem(int(chip.ADDR()))
    if chip.RI() == 1 : chip._dat(chip.DAT())
    return None

if __name__ == "__main__":
    mem = Memory(16, 8)
    mem.prog_mem(0, 0b00011110)
    mem.prog_mem(1, 0b01010101)
    mem.AI(1)
    mem.RO(1)
    mem.ADDR(1)
    for i in range(Memory.maxframe+1):
        mem.update(i)
        print(f"<frame>:[{i}]", mem)
